from text2vec import SentenceModel
import numpy as np
from sklearn.metrics.pairwise import cosine_similarity
import jieba  # 引入结巴分词
import time


class TrieNode:
    def __init__(self):
        self.children = {}  # 存储子节点的字典。键是字符，值是对应字符的 TrieNode 子节点。
        self.end_of_word = False  # 标记该节点是否为一个完整单词的结尾。

    def insert(self, word):
        """
        将单词插入到 Trie 中。

        :param word: 要插入的单词。
        """
        current = self
        for char in word:  # 遍历单词中的每一个字符
            if char not in current.children:
                current.children[char] = TrieNode()  # 如果当前字符不在子节点中，则创建一个新的 TrieNode 实例并添加到 children 字典中
            current = current.children[char]  # 移动到下一个字符的子节点
        current.end_of_word = True  # 在单词的最后一个字符节点上标记该节点为一个完整单词的结尾

    def search_prefix(self, prefix):
        """
        搜索前缀并返回所有以该前缀开始的单词。

        :param prefix: 要搜索的前缀字符串。
        :return: 以该前缀开始的所有单词列表。如果没有匹配的单词，则返回空列表。
        """
        current = self
        for char in prefix:  # 遍历前缀中的每一个字符
            if char not in current.children:
                return []  # 如果当前字符不在子节点中，则返回空列表，表示没有匹配的单词
            current = current.children[char]  # 移动到下一个字符的子节点
        return self._collect_words(current, prefix)  # 从当前节点开始收集所有以该前缀开始的单词

    def _collect_words(self, node, prefix):
        """
        从给定节点开始收集所有以当前前缀开始的单词。

        :param node: 当前处理的节点。
        :param prefix: 当前前缀字符串。
        :return: 收集到的所有单词列表。
        """
        words = []
        if node.end_of_word:
            words.append(prefix)  # 如果当前节点是一个完整单词的结尾，则将前缀添加到 words 列表中
        for char, next_node in node.children.items():  # 遍历当前节点的所有子节点
            words.extend(self._collect_words(next_node, prefix + char))  # 递归调用 _collect_words 方法，处理子节点并将结果扩展到 words 列表中
        return words  # 返回收集到的所有单词


def load_keywords(path):
    """
    从文件中加载关键词。

    :param path: 关键词文件的路径。
    :return: 关键词列表。
    """
    with open(path, 'r', encoding='utf-8') as file:
        keywords = [line.strip() for line in file.readlines()]  # 使用列表生成式从文件中读取并处理每一行
    return keywords


# 初始化 SentenceModel
model_path = "model/shibing624-text2vec-base-chinese"
model = SentenceModel(model_path)  # 导入并使用外部库中的 SentenceModel 类

# 从本地文件加载关键词
keywords_file_path = "dict/merged_dictionary.dat"
keywords_dict = load_keywords(keywords_file_path)  # 从文件中加载关键词

# 初始化 Trie 并插入关键词
trie = TrieNode()
for word in keywords_dict:
    trie.insert(word)  # 将关键词插入到 Trie 中

# 编码词典中的关键词
keyword_embeddings = model.encode(keywords_dict)

# 分词和编码查询
def tokenize_and_encode(query):
    tokens = jieba.lcut(query)
    token_embeddings = model.encode(tokens)
    return np.mean(token_embeddings, axis=0)

def find_similar_keywords(query, top_n=10):
    # 记录开始时间
    start_time = time.time()

    # 先进行前缀匹配
    prefix_matched_words = trie.search_prefix(query)

    # 如果前缀匹配结果不为空，只对这些结果计算相似度
    if prefix_matched_words:
        matched_embeddings = model.encode(prefix_matched_words)
    else:
        # 如果没有前缀匹配结果，退回到全词典搜索
        matched_embeddings = keyword_embeddings
        prefix_matched_words = keywords_dict

    # 计算相似度
    query_embedding = tokenize_and_encode(query)
    similarities = cosine_similarity([query_embedding], matched_embeddings)
    top_indices = np.argsort(similarities[0])[::-1][:top_n]
    similar_keywords = [(prefix_matched_words[i], similarities[0][i]) for i in top_indices]

    # 计算执行时间
    elapsed_time = time.time() - start_time
    return similar_keywords, elapsed_time

# 测试函数，添加循环允许多次查询
if __name__ == "__main__":
    while True:
        query = input("请输入查询词 (输入'退出'结束查询): ")
        if query.lower() == '退出':
            break
        similar_keywords, elapsed_time = find_similar_keywords(query)
        print(f"与 '{query}' 语义上接近的关键词和相应的相似度分数为:")
        for keyword, score in similar_keywords:
            print(f"{keyword}: {score:.4f}")
        print(f"查询耗时: {elapsed_time:.4f} 秒\n")
